package com.github.kezhenxu94.data.source;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.Iterator;
import java.util.NavigableMap;
import java.util.Objects;
import java.util.Properties;
import java.util.concurrent.*;
import java.util.function.Function;

/**
 * Created by kezhenxu on 1/5/17.
 *
 * @author kezhenxu (zhenxuke@newsjet-inc.com)
 */
public class HBaseSource<T> implements Source<T> {

    protected ExecutorService executor;
    protected Configuration configuration;
    protected Connection connection;
    protected Properties properties;
    protected BlockingQueue<T> queue;
    protected volatile boolean isRunning;
    protected String tableName;
    protected String columnFamily;
    protected Function<byte[], T> valueDeserializer;
    protected int batch = 1000;

    @Override
    public void setup() throws Exception {
        this.configuration = HBaseConfiguration.create();
        this.properties = new Properties();
        this.properties.load(getClass().getResourceAsStream("/hbase-source.properties"));
        this.properties.entrySet().forEach(property -> configuration.set(Objects.toString(property.getKey()),
                Objects.toString(property.getValue())));
        this.connection = ConnectionFactory.createConnection(configuration);
        this.queue = new LinkedBlockingQueue<>();
        this.executor = Executors.newSingleThreadExecutor();
    }

    @Override
    public void start() throws Exception {
        isRunning = true;
        executor.execute(() -> {
            try (Table table = connection.getTable(TableName.valueOf(tableName.getBytes()))) {
                Scan scan = new Scan();
                scan.setBatch(batch);
                try (ResultScanner results = table.getScanner(scan)) {
                    for (Iterator<Result> iterator = results.iterator(); iterator.hasNext() && isRunning; ) {
                        Result result = iterator.next();
                        NavigableMap<byte[], byte[]> familyMap = result.getFamilyMap(Bytes.toBytes(columnFamily));
                        familyMap.entrySet().forEach(entry -> queue.offer(valueDeserializer.apply(entry.getValue())));
                    }
                }
            } catch (IOException e) {
                getLogger().error(e.getMessage(), e);
            }
        });
    }

    @Override
    public void close() throws Exception {
        isRunning = false;
        try {
            while (isRunning) {
                TimeUnit.MILLISECONDS.sleep(500L);
            }
            executor.awaitTermination(5L, TimeUnit.SECONDS);
        } finally {
            if (connection != null) {
                connection.close();
            }
        }
    }

    @Override
    public boolean isRunning() {
        return isRunning || !queue.isEmpty();
    }

    @Override
    public T take() throws Exception {
        return isRunning() ? queue.take() : null;
    }

    public String getTableName() {
        return tableName;
    }

    public void setTableName(String tableName) {
        this.tableName = tableName;
    }

    public String getColumnFamily() {
        return columnFamily;
    }

    public void setColumnFamily(String columnFamily) {
        this.columnFamily = columnFamily;
    }

    public Function<byte[], T> getValueDeserializer() {
        return valueDeserializer;
    }

    public void setValueDeserializer(Function<byte[], T> valueDeserializer) {
        this.valueDeserializer = valueDeserializer;
    }

    public int getBatch() {
        return batch;
    }

    public void setBatch(int batch) {
        this.batch = batch;
    }
}
